// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2016
// Mehdi Goli    Codeplay Software Ltd.
// Ralph Potter  Codeplay Software Ltd.
// Luke Iwanski  Codeplay Software Ltd.
// Contact: <eigen@codeplay.com>
// Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#define EIGEN_TEST_NO_LONGDOUBLE
#define EIGEN_TEST_NO_COMPLEX

#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int64_t
#define EIGEN_USE_SYCL

#include "main.h"
#include <unsupported/Eigen/CXX11/Tensor>

using Eigen::array;
using Eigen::SyclDevice;
using Eigen::Tensor;
using Eigen::TensorMap;

template<typename DataType, int DataLayout, typename IndexType>
static void
test_simple_reshape(const Eigen::SyclDevice& sycl_device)
{
	typename Tensor<DataType, 5, DataLayout, IndexType>::Dimensions dim1(2, 3, 1, 7, 1);
	typename Tensor<DataType, 3, DataLayout, IndexType>::Dimensions dim2(2, 3, 7);
	typename Tensor<DataType, 2, DataLayout, IndexType>::Dimensions dim3(6, 7);
	typename Tensor<DataType, 2, DataLayout, IndexType>::Dimensions dim4(2, 21);

	Tensor<DataType, 5, DataLayout, IndexType> tensor1(dim1);
	Tensor<DataType, 3, DataLayout, IndexType> tensor2(dim2);
	Tensor<DataType, 2, DataLayout, IndexType> tensor3(dim3);
	Tensor<DataType, 2, DataLayout, IndexType> tensor4(dim4);

	tensor1.setRandom();

	DataType* gpu_data1 = static_cast<DataType*>(sycl_device.allocate(tensor1.size() * sizeof(DataType)));
	DataType* gpu_data2 = static_cast<DataType*>(sycl_device.allocate(tensor2.size() * sizeof(DataType)));
	DataType* gpu_data3 = static_cast<DataType*>(sycl_device.allocate(tensor3.size() * sizeof(DataType)));
	DataType* gpu_data4 = static_cast<DataType*>(sycl_device.allocate(tensor4.size() * sizeof(DataType)));

	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu1(gpu_data1, dim1);
	TensorMap<Tensor<DataType, 3, DataLayout, IndexType>> gpu2(gpu_data2, dim2);
	TensorMap<Tensor<DataType, 2, DataLayout, IndexType>> gpu3(gpu_data3, dim3);
	TensorMap<Tensor<DataType, 2, DataLayout, IndexType>> gpu4(gpu_data4, dim4);

	sycl_device.memcpyHostToDevice(gpu_data1, tensor1.data(), (tensor1.size()) * sizeof(DataType));

	gpu2.device(sycl_device) = gpu1.reshape(dim2);
	sycl_device.memcpyDeviceToHost(tensor2.data(), gpu_data2, (tensor1.size()) * sizeof(DataType));

	gpu3.device(sycl_device) = gpu1.reshape(dim3);
	sycl_device.memcpyDeviceToHost(tensor3.data(), gpu_data3, (tensor3.size()) * sizeof(DataType));

	gpu4.device(sycl_device) = gpu1.reshape(dim2).reshape(dim4);
	sycl_device.memcpyDeviceToHost(tensor4.data(), gpu_data4, (tensor4.size()) * sizeof(DataType));
	for (IndexType i = 0; i < 2; ++i) {
		for (IndexType j = 0; j < 3; ++j) {
			for (IndexType k = 0; k < 7; ++k) {
				VERIFY_IS_EQUAL(tensor1(i, j, 0, k, 0), tensor2(i, j, k)); /// ColMajor
				if (static_cast<int>(DataLayout) == static_cast<int>(ColMajor)) {
					VERIFY_IS_EQUAL(tensor1(i, j, 0, k, 0), tensor3(i + 2 * j, k)); /// ColMajor
					VERIFY_IS_EQUAL(tensor1(i, j, 0, k, 0), tensor4(i, j + 3 * k)); /// ColMajor
				} else {
					// VERIFY_IS_EQUAL(tensor1(i,j,0,k,0), tensor2(i,j,k));      /// RowMajor
					VERIFY_IS_EQUAL(tensor1(i, j, 0, k, 0), tensor4(i, j * 7 + k)); /// RowMajor
					VERIFY_IS_EQUAL(tensor1(i, j, 0, k, 0), tensor3(i * 3 + j, k)); /// RowMajor
				}
			}
		}
	}
	sycl_device.deallocate(gpu_data1);
	sycl_device.deallocate(gpu_data2);
	sycl_device.deallocate(gpu_data3);
	sycl_device.deallocate(gpu_data4);
}

template<typename DataType, int DataLayout, typename IndexType>
static void
test_reshape_as_lvalue(const Eigen::SyclDevice& sycl_device)
{
	typename Tensor<DataType, 3, DataLayout, IndexType>::Dimensions dim1(2, 3, 7);
	typename Tensor<DataType, 2, DataLayout, IndexType>::Dimensions dim2(6, 7);
	typename Tensor<DataType, 5, DataLayout, IndexType>::Dimensions dim3(2, 3, 1, 7, 1);
	Tensor<DataType, 3, DataLayout, IndexType> tensor(dim1);
	Tensor<DataType, 2, DataLayout, IndexType> tensor2d(dim2);
	Tensor<DataType, 5, DataLayout, IndexType> tensor5d(dim3);

	tensor.setRandom();

	DataType* gpu_data1 = static_cast<DataType*>(sycl_device.allocate(tensor.size() * sizeof(DataType)));
	DataType* gpu_data2 = static_cast<DataType*>(sycl_device.allocate(tensor2d.size() * sizeof(DataType)));
	DataType* gpu_data3 = static_cast<DataType*>(sycl_device.allocate(tensor5d.size() * sizeof(DataType)));

	TensorMap<Tensor<DataType, 3, DataLayout, IndexType>> gpu1(gpu_data1, dim1);
	TensorMap<Tensor<DataType, 2, DataLayout, IndexType>> gpu2(gpu_data2, dim2);
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu3(gpu_data3, dim3);

	sycl_device.memcpyHostToDevice(gpu_data1, tensor.data(), (tensor.size()) * sizeof(DataType));

	gpu2.reshape(dim1).device(sycl_device) = gpu1;
	sycl_device.memcpyDeviceToHost(tensor2d.data(), gpu_data2, (tensor2d.size()) * sizeof(DataType));

	gpu3.reshape(dim1).device(sycl_device) = gpu1;
	sycl_device.memcpyDeviceToHost(tensor5d.data(), gpu_data3, (tensor5d.size()) * sizeof(DataType));

	for (IndexType i = 0; i < 2; ++i) {
		for (IndexType j = 0; j < 3; ++j) {
			for (IndexType k = 0; k < 7; ++k) {
				VERIFY_IS_EQUAL(tensor5d(i, j, 0, k, 0), tensor(i, j, k));
				if (static_cast<int>(DataLayout) == static_cast<int>(ColMajor)) {
					VERIFY_IS_EQUAL(tensor2d(i + 2 * j, k), tensor(i, j, k)); /// ColMajor
				} else {
					VERIFY_IS_EQUAL(tensor2d(i * 3 + j, k), tensor(i, j, k)); /// RowMajor
				}
			}
		}
	}
	sycl_device.deallocate(gpu_data1);
	sycl_device.deallocate(gpu_data2);
	sycl_device.deallocate(gpu_data3);
}

template<typename DataType, int DataLayout, typename IndexType>
static void
test_simple_slice(const Eigen::SyclDevice& sycl_device)
{
	IndexType sizeDim1 = 2;
	IndexType sizeDim2 = 3;
	IndexType sizeDim3 = 5;
	IndexType sizeDim4 = 7;
	IndexType sizeDim5 = 11;
	array<IndexType, 5> tensorRange = { { sizeDim1, sizeDim2, sizeDim3, sizeDim4, sizeDim5 } };
	Tensor<DataType, 5, DataLayout, IndexType> tensor(tensorRange);
	tensor.setRandom();
	array<IndexType, 5> slice1_range = { { 1, 1, 1, 1, 1 } };
	Tensor<DataType, 5, DataLayout, IndexType> slice1(slice1_range);

	DataType* gpu_data1 = static_cast<DataType*>(sycl_device.allocate(tensor.size() * sizeof(DataType)));
	DataType* gpu_data2 = static_cast<DataType*>(sycl_device.allocate(slice1.size() * sizeof(DataType)));
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu1(gpu_data1, tensorRange);
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu2(gpu_data2, slice1_range);
	Eigen::DSizes<IndexType, 5> indices(1, 2, 3, 4, 5);
	Eigen::DSizes<IndexType, 5> sizes(1, 1, 1, 1, 1);
	sycl_device.memcpyHostToDevice(gpu_data1, tensor.data(), (tensor.size()) * sizeof(DataType));
	gpu2.device(sycl_device) = gpu1.slice(indices, sizes);
	sycl_device.memcpyDeviceToHost(slice1.data(), gpu_data2, (slice1.size()) * sizeof(DataType));
	VERIFY_IS_EQUAL(slice1(0, 0, 0, 0, 0), tensor(1, 2, 3, 4, 5));

	array<IndexType, 5> slice2_range = { { 1, 1, 2, 2, 3 } };
	Tensor<DataType, 5, DataLayout, IndexType> slice2(slice2_range);
	DataType* gpu_data3 = static_cast<DataType*>(sycl_device.allocate(slice2.size() * sizeof(DataType)));
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu3(gpu_data3, slice2_range);
	Eigen::DSizes<IndexType, 5> indices2(1, 1, 3, 4, 5);
	Eigen::DSizes<IndexType, 5> sizes2(1, 1, 2, 2, 3);
	gpu3.device(sycl_device) = gpu1.slice(indices2, sizes2);
	sycl_device.memcpyDeviceToHost(slice2.data(), gpu_data3, (slice2.size()) * sizeof(DataType));
	for (IndexType i = 0; i < 2; ++i) {
		for (IndexType j = 0; j < 2; ++j) {
			for (IndexType k = 0; k < 3; ++k) {
				VERIFY_IS_EQUAL(slice2(0, 0, i, j, k), tensor(1, 1, 3 + i, 4 + j, 5 + k));
			}
		}
	}
	sycl_device.deallocate(gpu_data1);
	sycl_device.deallocate(gpu_data2);
	sycl_device.deallocate(gpu_data3);
}

template<typename DataType, int DataLayout, typename IndexType>
static void
test_strided_slice_as_rhs_sycl(const Eigen::SyclDevice& sycl_device)
{
	IndexType sizeDim1 = 2;
	IndexType sizeDim2 = 3;
	IndexType sizeDim3 = 5;
	IndexType sizeDim4 = 7;
	IndexType sizeDim5 = 11;
	typedef Eigen::DSizes<IndexType, 5> Index5;
	Index5 strides(1L, 1L, 1L, 1L, 1L);
	Index5 indicesStart(1L, 2L, 3L, 4L, 5L);
	Index5 indicesStop(2L, 3L, 4L, 5L, 6L);
	Index5 lengths(1L, 1L, 1L, 1L, 1L);

	array<IndexType, 5> tensorRange = { { sizeDim1, sizeDim2, sizeDim3, sizeDim4, sizeDim5 } };
	Tensor<DataType, 5, DataLayout, IndexType> tensor(tensorRange);
	tensor.setRandom();

	array<IndexType, 5> slice1_range = { { 1, 1, 1, 1, 1 } };
	Tensor<DataType, 5, DataLayout, IndexType> slice1(slice1_range);
	Tensor<DataType, 5, DataLayout, IndexType> slice_stride1(slice1_range);

	DataType* gpu_data1 = static_cast<DataType*>(sycl_device.allocate(tensor.size() * sizeof(DataType)));
	DataType* gpu_data2 = static_cast<DataType*>(sycl_device.allocate(slice1.size() * sizeof(DataType)));
	DataType* gpu_data_stride2 = static_cast<DataType*>(sycl_device.allocate(slice_stride1.size() * sizeof(DataType)));

	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu1(gpu_data1, tensorRange);
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu2(gpu_data2, slice1_range);
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu_stride2(gpu_data_stride2, slice1_range);

	Eigen::DSizes<IndexType, 5> indices(1, 2, 3, 4, 5);
	Eigen::DSizes<IndexType, 5> sizes(1, 1, 1, 1, 1);
	sycl_device.memcpyHostToDevice(gpu_data1, tensor.data(), (tensor.size()) * sizeof(DataType));
	gpu2.device(sycl_device) = gpu1.slice(indices, sizes);
	sycl_device.memcpyDeviceToHost(slice1.data(), gpu_data2, (slice1.size()) * sizeof(DataType));

	gpu_stride2.device(sycl_device) = gpu1.stridedSlice(indicesStart, indicesStop, strides);
	sycl_device.memcpyDeviceToHost(slice_stride1.data(), gpu_data_stride2, (slice_stride1.size()) * sizeof(DataType));

	VERIFY_IS_EQUAL(slice1(0, 0, 0, 0, 0), tensor(1, 2, 3, 4, 5));
	VERIFY_IS_EQUAL(slice_stride1(0, 0, 0, 0, 0), tensor(1, 2, 3, 4, 5));

	array<IndexType, 5> slice2_range = { { 1, 1, 2, 2, 3 } };
	Tensor<DataType, 5, DataLayout, IndexType> slice2(slice2_range);
	Tensor<DataType, 5, DataLayout, IndexType> strideSlice2(slice2_range);

	DataType* gpu_data3 = static_cast<DataType*>(sycl_device.allocate(slice2.size() * sizeof(DataType)));
	DataType* gpu_data_stride3 = static_cast<DataType*>(sycl_device.allocate(strideSlice2.size() * sizeof(DataType)));
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu3(gpu_data3, slice2_range);
	TensorMap<Tensor<DataType, 5, DataLayout, IndexType>> gpu_stride3(gpu_data_stride3, slice2_range);
	Eigen::DSizes<IndexType, 5> indices2(1, 1, 3, 4, 5);
	Eigen::DSizes<IndexType, 5> sizes2(1, 1, 2, 2, 3);
	Index5 strides2(1L, 1L, 1L, 1L, 1L);
	Index5 indicesStart2(1L, 1L, 3L, 4L, 5L);
	Index5 indicesStop2(2L, 2L, 5L, 6L, 8L);

	gpu3.device(sycl_device) = gpu1.slice(indices2, sizes2);
	sycl_device.memcpyDeviceToHost(slice2.data(), gpu_data3, (slice2.size()) * sizeof(DataType));

	gpu_stride3.device(sycl_device) = gpu1.stridedSlice(indicesStart2, indicesStop2, strides2);
	sycl_device.memcpyDeviceToHost(strideSlice2.data(), gpu_data_stride3, (strideSlice2.size()) * sizeof(DataType));

	for (IndexType i = 0; i < 2; ++i) {
		for (IndexType j = 0; j < 2; ++j) {
			for (IndexType k = 0; k < 3; ++k) {
				VERIFY_IS_EQUAL(slice2(0, 0, i, j, k), tensor(1, 1, 3 + i, 4 + j, 5 + k));
				VERIFY_IS_EQUAL(strideSlice2(0, 0, i, j, k), tensor(1, 1, 3 + i, 4 + j, 5 + k));
			}
		}
	}
	sycl_device.deallocate(gpu_data1);
	sycl_device.deallocate(gpu_data2);
	sycl_device.deallocate(gpu_data3);
}

template<typename DataType, int DataLayout, typename IndexType>
static void
test_strided_slice_write_sycl(const Eigen::SyclDevice& sycl_device)
{
	typedef Tensor<DataType, 2, DataLayout, IndexType> Tensor2f;
	typedef Eigen::DSizes<IndexType, 2> Index2;
	IndexType sizeDim1 = 7L;
	IndexType sizeDim2 = 11L;
	array<IndexType, 2> tensorRange = { { sizeDim1, sizeDim2 } };
	Tensor<DataType, 2, DataLayout, IndexType> tensor(tensorRange), tensor2(tensorRange);
	IndexType sliceDim1 = 2;
	IndexType sliceDim2 = 3;
	array<IndexType, 2> sliceRange = { { sliceDim1, sliceDim2 } };
	Tensor2f slice(sliceRange);
	Index2 strides(1L, 1L);
	Index2 indicesStart(3L, 4L);
	Index2 indicesStop(5L, 7L);
	Index2 lengths(2L, 3L);

	DataType* gpu_data1 = static_cast<DataType*>(sycl_device.allocate(tensor.size() * sizeof(DataType)));
	DataType* gpu_data2 = static_cast<DataType*>(sycl_device.allocate(tensor2.size() * sizeof(DataType)));
	DataType* gpu_data3 = static_cast<DataType*>(sycl_device.allocate(slice.size() * sizeof(DataType)));
	TensorMap<Tensor<DataType, 2, DataLayout, IndexType>> gpu1(gpu_data1, tensorRange);
	TensorMap<Tensor<DataType, 2, DataLayout, IndexType>> gpu2(gpu_data2, tensorRange);
	TensorMap<Tensor<DataType, 2, DataLayout, IndexType>> gpu3(gpu_data3, sliceRange);

	tensor.setRandom();
	sycl_device.memcpyHostToDevice(gpu_data1, tensor.data(), (tensor.size()) * sizeof(DataType));
	gpu2.device(sycl_device) = gpu1;

	slice.setRandom();
	sycl_device.memcpyHostToDevice(gpu_data3, slice.data(), (slice.size()) * sizeof(DataType));

	gpu1.slice(indicesStart, lengths).device(sycl_device) = gpu3;
	gpu2.stridedSlice(indicesStart, indicesStop, strides).device(sycl_device) = gpu3;
	sycl_device.memcpyDeviceToHost(tensor.data(), gpu_data1, (tensor.size()) * sizeof(DataType));
	sycl_device.memcpyDeviceToHost(tensor2.data(), gpu_data2, (tensor2.size()) * sizeof(DataType));

	for (IndexType i = 0; i < sizeDim1; i++)
		for (IndexType j = 0; j < sizeDim2; j++) {
			VERIFY_IS_EQUAL(tensor(i, j), tensor2(i, j));
		}
	sycl_device.deallocate(gpu_data1);
	sycl_device.deallocate(gpu_data2);
	sycl_device.deallocate(gpu_data3);
}

template<typename OutIndex, typename DSizes>
Eigen::array<OutIndex, DSizes::count>
To32BitDims(const DSizes& in)
{
	Eigen::array<OutIndex, DSizes::count> out;
	for (int i = 0; i < DSizes::count; ++i) {
		out[i] = in[i];
	}
	return out;
}

template<class DataType, int DataLayout, typename IndexType, typename ConvertedIndexType>
int
run_eigen(const SyclDevice& sycl_device)
{
	using TensorI64 = Tensor<DataType, 5, DataLayout, IndexType>;
	using TensorI32 = Tensor<DataType, 5, DataLayout, ConvertedIndexType>;
	using TensorMI64 = TensorMap<TensorI64>;
	using TensorMI32 = TensorMap<TensorI32>;
	Eigen::array<IndexType, 5> tensor_range{ { 4, 1, 1, 1, 6 } };
	Eigen::array<IndexType, 5> slice_range{ { 4, 1, 1, 1, 3 } };

	TensorI64 out_tensor_gpu(tensor_range);
	TensorI64 out_tensor_cpu(tensor_range);
	out_tensor_cpu.setRandom();

	TensorI64 sub_tensor(slice_range);
	sub_tensor.setRandom();

	DataType* out_gpu_data = static_cast<DataType*>(sycl_device.allocate(out_tensor_cpu.size() * sizeof(DataType)));
	DataType* sub_gpu_data = static_cast<DataType*>(sycl_device.allocate(sub_tensor.size() * sizeof(DataType)));
	TensorMI64 out_gpu(out_gpu_data, tensor_range);
	TensorMI64 sub_gpu(sub_gpu_data, slice_range);

	sycl_device.memcpyHostToDevice(out_gpu_data, out_tensor_cpu.data(), out_tensor_cpu.size() * sizeof(DataType));
	sycl_device.memcpyHostToDevice(sub_gpu_data, sub_tensor.data(), sub_tensor.size() * sizeof(DataType));

	Eigen::array<ConvertedIndexType, 5> slice_offset_32{ { 0, 0, 0, 0, 3 } };
	Eigen::array<ConvertedIndexType, 5> slice_range_32{ { 4, 1, 1, 1, 3 } };
	TensorMI32 out_cpu_32(out_tensor_cpu.data(), To32BitDims<ConvertedIndexType>(out_tensor_cpu.dimensions()));
	TensorMI32 sub_cpu_32(sub_tensor.data(), To32BitDims<ConvertedIndexType>(sub_tensor.dimensions()));
	TensorMI32 out_gpu_32(out_gpu.data(), To32BitDims<ConvertedIndexType>(out_gpu.dimensions()));
	TensorMI32 sub_gpu_32(sub_gpu.data(), To32BitDims<ConvertedIndexType>(sub_gpu.dimensions()));

	out_gpu_32.slice(slice_offset_32, slice_range_32).device(sycl_device) = sub_gpu_32;

	out_cpu_32.slice(slice_offset_32, slice_range_32) = sub_cpu_32;

	sycl_device.memcpyDeviceToHost(out_tensor_gpu.data(), out_gpu_data, out_tensor_cpu.size() * sizeof(DataType));
	int has_err = 0;
	for (IndexType i = 0; i < out_tensor_cpu.size(); ++i) {
		auto exp = out_tensor_cpu(i);
		auto val = out_tensor_gpu(i);
		if (val != exp) {
			std::cout << "#" << i << " got " << val << " but expected " << exp << std::endl;
			has_err = 1;
		}
	}
	sycl_device.deallocate(out_gpu_data);
	sycl_device.deallocate(sub_gpu_data);
	return has_err;
}

template<typename DataType, typename dev_Selector>
void
sycl_morphing_test_per_device(dev_Selector s)
{
	QueueInterface queueInterface(s);
	auto sycl_device = Eigen::SyclDevice(&queueInterface);
	test_simple_slice<DataType, RowMajor, int64_t>(sycl_device);
	test_simple_slice<DataType, ColMajor, int64_t>(sycl_device);
	test_simple_reshape<DataType, RowMajor, int64_t>(sycl_device);
	test_simple_reshape<DataType, ColMajor, int64_t>(sycl_device);
	test_reshape_as_lvalue<DataType, RowMajor, int64_t>(sycl_device);
	test_reshape_as_lvalue<DataType, ColMajor, int64_t>(sycl_device);
	test_strided_slice_write_sycl<DataType, ColMajor, int64_t>(sycl_device);
	test_strided_slice_write_sycl<DataType, RowMajor, int64_t>(sycl_device);
	test_strided_slice_as_rhs_sycl<DataType, ColMajor, int64_t>(sycl_device);
	test_strided_slice_as_rhs_sycl<DataType, RowMajor, int64_t>(sycl_device);
	run_eigen<float, RowMajor, long, int>(sycl_device);
}
EIGEN_DECLARE_TEST(cxx11_tensor_morphing_sycl)
{
	for (const auto& device : Eigen::get_sycl_supported_devices()) {
		CALL_SUBTEST(sycl_morphing_test_per_device<float>(device));
	}
}
